package com.gallop.mq.core.template;

import com.gallop.mq.MQProperties;
import com.gallop.mq.MQTemplate;
import com.gallop.mq.core.MQExecutorProvider;
import com.gallop.mq.core.MQMessageCodec;
import com.gallop.mq.core.MQMessageDispatcher;
import com.gallop.mq.core.dispatcher.DispatchContext;
import com.gallop.mq.core.exception.MQException;
import com.gallop.mq.core.template.handler.DeliveryContext;
import com.gallop.mq.core.template.handler.DeliveryType;
import com.gallop.mq.core.template.handler.MessageDeliveryHandler;
import com.gallop.mq.core.template.handler.registry.DynamicProxyListenerCreator;
import com.gallop.mq.core.template.handler.registry.MessageDeliveryHandlerRegistry;
import com.gallop.mq.core.template.handler.registry.RegistryHolder;
import com.gallop.mq.core.template.handler.registry.RegistryPreprocessor;
import com.google.common.collect.Maps;
import com.rabbitmq.client.*;
import lombok.EqualsAndHashCode;
import lombok.Getter;
import lombok.Setter;
import lombok.ToString;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.StringUtils;

import java.io.IOException;
import java.time.Duration;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

@Getter
@Setter
@Slf4j
public class RabbitMQTemplate implements MQTemplate, RegistryHolder {

    private MQMessageCodec codec;
    private MQMessageDispatcher dispatcher;

    private String localReplyQueue;

    private Connection connection;

    private Channel channel;

    private MQProperties properties;

    private Map<DynamicExchange, PendingReply> pendingReplyMap = Maps.newConcurrentMap();

    private MessageDeliveryHandlerRegistry handlerRegistry;

    private void internalPublish(String topic, String queue, AMQP.BasicProperties properties, Object message) throws Exception {
        try (Channel ch = connection.createChannel()) {
            String exchangeName = "temp-exchange-" + UUID.randomUUID().toString().replaceAll("-", "");
            ch.exchangeDeclare(exchangeName, BuiltinExchangeType.DIRECT, false, false, null);
            ch.queueDeclare(queue, true, false, true, null);
            ch.queueBind(queue, exchangeName, topic);
            ch.basicPublish(exchangeName, topic, properties, codec.encode(message));

            ch.queueUnbind(queue, exchangeName, topic);
            ch.exchangeDelete(exchangeName);
        }
    }

    @Override
    public void init(MQProperties properties, MQMessageDispatcher dispatcher, MQMessageCodec codec, MQExecutorProvider executorProvider, MessageDeliveryHandlerRegistry handlerRegistry) throws MQException {
        setDispatcher(dispatcher);
        setCodec(codec);
        setProperties(properties);

        this.initRegistry(handlerRegistry);

        try {
            prepareLocalEntry(properties, executorProvider);
            prepareTopicListeners();
            log.info("MQ RabbitMQ init completed. connect to: {},RpcQueue: {}, [{}] topic(s) scanned in total.",
                    properties.getHosts(), localReplyQueue, this.dispatcher.getAllContainers().size());
        } catch (Exception e) {
            log.error("", e);
            throw new MQException("", e);
        }
    }

    private void prepareLocalEntry(MQProperties properties, MQExecutorProvider executorProvider) throws Exception {
        connection = createRabbitMQConnection(properties, executorProvider);
        channel = connection.createChannel();
        if (channel == null) {
            throw new MQException("init mq template failed. property=" + properties);
        }
        // 本地发送消息exchange
        String localName = UUID.randomUUID().toString().replaceAll("-", "");
        // 本地创建rpc响应队列并监听
        localReplyQueue = properties.getRabbitmq().getRpcReplyQueue();
        if (!StringUtils.hasText(localReplyQueue)) {
            localReplyQueue = "local-reply-" + localName;
        }
        this.channel.queueDeclare(localReplyQueue, true, false, true, null);
//        this.channel.basicConsume(localReplyQueue, true, DefaultMessageHandlerAdaptor.of(this));
        this.channel.basicConsume(localReplyQueue, true,
                handlerRegistry.createListener(DeliveryContext.builder().mqTemplate(this).build()));
    }

    private Connection createRabbitMQConnection(MQProperties properties,
                                                MQExecutorProvider executorProvider) throws Exception {
        ConnectionFactory factory = new ConnectionFactory();
        factory.setUsername(properties.getUser());
        factory.setPassword(properties.getSecret());
        String uriString = properties.getHosts().get(0);
        if (!uriString.startsWith("amqp") && !uriString.startsWith("amqps")) {
            uriString = "amqp://" + uriString;
        }
        factory.setUri(uriString);
        factory.setVirtualHost(properties.getRabbitmq().getVirtualHost());
        factory.setAutomaticRecoveryEnabled(true);
        factory.setTopologyRecoveryEnabled(true);
        factory.setSharedExecutor(executorProvider.create());
        factory.setNetworkRecoveryInterval(properties.getReconnectInterval());
        return factory.newConnection();
    }

    private void prepareTopicListeners() {
        // 监听扫描到的topic
        this.dispatcher.getAllContainers().forEach(container -> {
            String topic = container.getTopic();
            String group = container.getGroup();// TODO group
            try {
                // 监听消息
                channel.queueDeclare(topic, true, false, true, null);
//                channel.basicConsume(topic, true, DefaultMessageHandlerAdaptor.of(this, topic, group));
                channel.basicConsume(topic, true, handlerRegistry.createListener(DeliveryContext.builder()
                        .mqTemplate(this)
                        .container(container)
                        .build()));
            } catch (Exception e) {
                throw new MQException("queue consumer declaring failed.topic=" + topic, e);
            }
        });
    }

    @Override
    public void onDispose() {
        try {
            if (channel != null && channel.isOpen()) {
                log.info("MQTemplate RabbitMQ channel destroy.");
                pendingReplyMap.forEach((exchange, reply) -> {
                    try {
                        channel.queueUnbind(reply.getSavedReplyTo(), exchange.getId(), reply.getSavedCorrelation());
                        channel.exchangeDelete(exchange.getId());
                    } catch (IOException e) {
                        log.error("error when shutdown exchange", e);
                    }
                });
                pendingReplyMap.clear();
                channel.queueDelete(localReplyQueue);
                channel.close();
            }
            if (connection != null && connection.isOpen()) {
                log.info("MQTemplate RabbitMQ connection destroy.");
                connection.close();
            }
        } catch (Exception e) {
            log.error("", e);
        }
    }


    @Override
    public void send(Object message, String topic) {
        try {
            internalPublish(topic, topic, new AMQP.BasicProperties.Builder()
                    .type(DeliveryType.SEND.name())
                    .build(), message);
        } catch (Exception e) {
            throw new MQException("send message failed.", e);
        }
    }

    @Override
    public <P> void request(Object message, String topic, Duration requestTimeout, Consumer<P> onResponse) {
        request(message, topic, requestTimeout, onResponse, (msg, ex) -> log.error("mq request failed. msg=" + msg, ex));
    }

    @Override
    public <P> void request(Object message, String topic, Duration requestTimeout, Consumer<P> onResponse, BiConsumer<Object, Exception> onError) {
        PendingReply pendingReply = new PendingReply();
        pendingReply.setSavedReplyTo(localReplyQueue);
        String savedCorrelation = "rpc-exchange-" + UUID.randomUUID().toString().replaceAll("-", "");
        pendingReply.setSavedCorrelation(savedCorrelation);
        DynamicExchange dynamicExchange = new DynamicExchange(savedCorrelation);
        pendingReplyMap.put(dynamicExchange, pendingReply);
        try {
            // send
            AMQP.BasicProperties prop = new AMQP.BasicProperties.Builder()
                    .correlationId(savedCorrelation)
                    .replyTo(localReplyQueue)
                    .type(DeliveryType.REQUEST.name())
                    .build();
            internalPublish(topic, topic, prop, message);
            P response;
            if (requestTimeout == null && properties.getRequestTimeout() > 0) {
                requestTimeout = Duration.ofMillis(properties.getRequestTimeout());
            }
            if (requestTimeout == null) {
                response = pendingReply.get();
            } else {
                response = pendingReply.get(requestTimeout.toMillis(), TimeUnit.MILLISECONDS);
            }
            onResponse.accept(response);
        } catch (Exception e) {
            pendingReplyMap.remove(dynamicExchange);
            pendingReply.returned(e);
            if (onError != null) {
                onError.accept(null, e);
            }
        }
    }

    @Override
    public <P> P request(Object message, String topic) throws InterruptedException {
        Integer requestTimeout = this.properties.getRequestTimeout();
        if (requestTimeout == null || requestTimeout <= 0) {
            return request(message, topic, (Duration) null);
        }
        return request(message, topic, Duration.ofMillis(requestTimeout));
    }

    @Override
    public <P> P request(Object message, String topic, Duration receiveTimeout) throws InterruptedException {
        PendingReply pendingReply = new PendingReply();
        pendingReply.setSavedReplyTo(localReplyQueue);
        String savedCorrelation = "rpc-exchange-" + UUID.randomUUID().toString().replaceAll("-", "");
        pendingReply.setSavedCorrelation(savedCorrelation);
        DynamicExchange dynamicExchange = new DynamicExchange(savedCorrelation);
        pendingReplyMap.put(dynamicExchange, pendingReply);
        try {
            // send
            AMQP.BasicProperties prop = new AMQP.BasicProperties.Builder()
                    .correlationId(pendingReply.getSavedCorrelation())
                    .replyTo(localReplyQueue)
                    .type(DeliveryType.REQUEST.name())
                    .build();
            internalPublish(topic, topic, prop, message);
            P response;
            if (receiveTimeout == null) {
                response = pendingReply.get();
            } else {
                response = pendingReply.get(receiveTimeout.toMillis(), TimeUnit.MILLISECONDS);
            }
            return response;
        } catch (Exception e) {
            log.error("mq request failed.", e);
            pendingReplyMap.remove(dynamicExchange);
            pendingReply.returned(e);
        }
        return null;
    }

    @Override
    public MessageDeliveryHandlerRegistry.InitializerBuilder registryInitializer() {
        return MessageDeliveryHandlerRegistry.InitializerBuilder.newBuilder()
                .setPreprocessor(RegistryPreprocessor.<com.rabbitmq.client.Consumer, byte[]>builder()
                        .listenerCreator(DynamicProxyListenerCreator.ofInterface(handlerRegistry, com.rabbitmq.client.Consumer.class))
                        .deliveryTypeProvider(args -> {
                            AMQP.BasicProperties prop = args.getArgAs(2);
                            return DeliveryType.valueOf(prop.getType());
                        })
                        .rpcInfoProvider(args -> {
                            AMQP.BasicProperties prop = args.getArgAs(2);
                            return MessageDeliveryHandler.StringRpcInfo.builder()
                                    .replyTo(prop.getReplyTo())
                                    .correlationId(prop.getCorrelationId())
                                    .build();
                        })
                        .rawMessageProvider(args -> args.getArgAs(3))
                        .mqMessageProvider(codec::decode).build())
                .registerHandler(DeliveryType.SEND, (deliveryContext, deliveryType, rpcInfo, incoming) ->
                        dispatcher.dispatch(DispatchContext.builder()
                                .topic(deliveryContext.getContainer().getTopic())
                                .group(deliveryContext.getContainer().getGroup())
                                .message(incoming)
                                .build()))
                .registerHandler(DeliveryType.REQUEST, (deliveryContext, deliveryType, rpcInfo, incoming) -> {
                    // 收到需要回复的rpc请求
                    dispatcher.dispatch(DispatchContext.builder()
                            .topic(deliveryContext.getContainer().getTopic())
                            .group(deliveryContext.getContainer().getGroup())
                            .message(incoming)
                            .replyTo((String) rpcInfo.getReplyTo())
                            .correlationId(rpcInfo.getCorrelationId())
                            .rpcResponse((correlationId, replyTo, msg) -> {
                                try {
                                    internalPublish(correlationId, replyTo, new AMQP.BasicProperties.Builder()
                                            .replyTo(replyTo)
                                            .correlationId(correlationId)
                                            .type(DeliveryType.RESPONSE.name())
                                            .build(), msg);
                                } catch (Exception e) {
                                    log.error("", e);
                                }
                            }).build());
                })
                .registerHandler(DeliveryType.RESPONSE, (deliveryContext, deliveryType, rpcInfo, incoming) -> {
                    PendingReply pendingReply = pendingReplyMap.remove(new DynamicExchange(rpcInfo.getCorrelationId()));
                    if (pendingReply != null) {
                        pendingReply.reply(incoming);
                    } else {
                        log.error("MQ response will be ignored. correlation={},replyTo={},message:{},",
                                rpcInfo.getCorrelationId(), rpcInfo.getReplyTo(), incoming);
                    }
                });
    }

    @Getter
    @Setter
    @ToString
    @EqualsAndHashCode
    private static class DynamicExchange {
        public DynamicExchange(String id) {
            this.id = id;
        }

        private String id;
    }
}
